################################################################################
#
# Copyright 2025 ByteDance Ltd. and/or its affiliates. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
################################################################################

import ctypes
import functools
import importlib
import logging
import os
from pathlib import Path

import torch

FLUX_TORCH_EXTENSION_NAME = "flux_ths_pybind"


def _preload_libs(libname, extra_path=None):
    libdirs = [Path(__file__).parent / "lib"]
    if extra_path is not None:
        libdirs.append(extra_path)

    for libdir in libdirs:
        try:
            ctypes.CDLL(libdir / libname)
            return
        except OSError:
            pass

    logging.debug(f"failed to load {libname} from {libdirs}. try load from LD_LIBRARY_PATH")
    ctypes.CDLL(libname)


def _nvshmem_home_by_pip():
    try:
        import nvidia.nvshmem

        return nvidia.nvshmem.__path__[0]
    except Exception:
        raise Exception(
            f"please install nvshmem by `pip3 install nvidia-nvshmem-cu12` or set NVSHMEM_HOME to an existing NVSHMEM"
        )


@functools.lru_cache()
def _nvshmem_home():
    return os.getenv("NVSHMEM_HOME", _nvshmem_home_by_pip())


def _load_deps():
    try:
        _preload_libs("libnvshmem_host.so.3", Path(_nvshmem_home()) / "lib")
        _preload_libs("nvshmem_bootstrap_uid.so.3", Path(_nvshmem_home()) / "lib")
        _preload_libs("nvshmem_transport_ibrc.so.3", Path(_nvshmem_home()) / "lib")
    except Exception as e:
        logging.warning("Failed to load NVSHMEM libs")
    _preload_libs("libflux_cuda.so")
    try:
        _preload_libs("libflux_triton_aot.so")
    except Exception as e:
        logging.debug("Failed to load triton_aot libs")
    _preload_libs("libflux_cuda_ths_op.so")


_load_deps()
flux_mod = importlib.import_module(FLUX_TORCH_EXTENSION_NAME)


class NotCompiled:
    pass


def _get_flux_member(member):
    return getattr(flux_mod, member, NotCompiled())


def _get_from_torch_classes(name: str):
    try:
        return getattr(torch.classes.flux, name)
    except:
        return NotCompiled()


bsr_reduce = _get_flux_member("bsr_reduce")
init_flux_shm = flux_mod.init_flux_shm
bitwise_check = flux_mod.bitwise_check
uniform_initialize = flux_mod.uniform_initialize
load_tuning_record = flux_mod.load_tuning_record
create_tensor_list = flux_mod.flux_create_tensor_list

_nvshmem_team_split_strided = _get_flux_member("_nvshmem_team_split_strided")
_nvshmem_team_split_2d = _get_flux_member("_nvshmem_team_split_2d")
_nvshmem_team_destroy = _get_flux_member("_nvshmem_team_destroy")
_nvshmem_team_get_config = _get_flux_member("_nvshmem_team_get_config")
_nvshmem_team_translate_pe = _get_flux_member("_nvshmem_team_translate_pe")
_nvshmem_team_npes = _get_flux_member("_nvshmem_team_npes")
_nvshmem_team_my_pe = _get_flux_member("_nvshmem_team_my_pe")

GroupBarrier = flux_mod.GroupBarrier

calc_scatter_index = _get_flux_member("calc_scatter_index")

ProfilingContext = flux_mod.ProfilingContext
TuningRecord = flux_mod.TuningRecord
DistEnvTP = flux_mod.DistEnvTP
DistEnvTPWithEP = flux_mod.DistEnvTPWithEP
MoeArguments = flux_mod.MoeArguments

# GEMM only
GemmOnly = _get_flux_member("GemmOnly")
BlockScaleGemm = _get_flux_member("BlockScaleGemm")
GemmGroupedV2 = _get_flux_member("GemmGroupedV2")
GemmGroupedV3 = _get_flux_member("GemmGroupedV3")

# GEMM+RS
GemmRS = _get_flux_member("GemmRS")
get_gemm_rs_threadblock_segments_info = _get_flux_member("get_gemm_rs_threadblock_segments_info")
calc_gemm_rs_threadblock_segments_info = _get_flux_member("calc_gemm_rs_threadblock_segments_info")
GemmRSInterNode = _get_flux_member("GemmRSInterNode")

# allgather op
AGRingMode = _get_flux_member("AGRingMode")
AllGatherOption = _get_flux_member("AllGatherOption")
if not isinstance(AllGatherOption, NotCompiled):
    AllGatherOption.__repr__ = (
        lambda x: f"AllGatherOption(mode={x.mode}, fuse_sync={x.fuse_sync}, use_cuda_core_local={x.use_cuda_core_local}, use_cuda_core_ag={x.use_cuda_core_ag}, input_buffer_copied={x.input_buffer_copied}, use_read={x.use_read})"
    )
AllGatherOp = _get_flux_member("AllGatherOp")
RingMode = _get_flux_member("RingMode")
ReduceScatterOption = _get_flux_member("ReduceScatterOption")
if not isinstance(ReduceScatterOption, NotCompiled):
    ReduceScatterOption.__repr__ = (
        lambda x: f"use_barrier_queue={x.use_barrier_queue}, "
        "use_1d_ring={x.use_1d_ring}, "
        "use_p2p_read={x.use_p2p_read}, "
        "use_cudaMemcpyAsync={x.use_cudaMemcpyAsync}, "
        "use_gemmk={x.use_gemmk}, "
        "per_tile_flags={x.per_tile_flags}, "
        "n_split={x.n_split}, "
        "num_blocks={x.num_blocks}, "
        "ring_mode={x.ring_mode}"
    )

# AG+GEMM
AGKernel = _get_flux_member("AGKernel")
AGKernelInterNode = _get_flux_member("AGKernelInterNode")

# MOE ag-scatter
GemmGroupedV2AGScatterOp = _get_flux_member("GemmGroupedV2AGScatterOp")
GemmGroupedV3AGScatter = _get_flux_member("GemmGroupedV3AGScatter")
prepare_moe_ag_scatter_args = _get_flux_member("prepare_moe_ag_scatter_args")

# MOE gather-rs
GemmGroupedV2GatherRSOp = _get_flux_member("GemmGroupedV2GatherRSOp")
TopkReduceScatterOp = _get_flux_member("TopkReduceScatterOp")
calc_moe_triton_blocked_gather_a = getattr(
    flux_mod, "calc_moe_triton_blocked_gather_a", NotCompiled()
)
GemmGroupedV3GatherRS = _get_flux_member("GemmGroupedV3GatherRS")
topk_scatter_reduce = _get_flux_member("topk_scatter_reduce")
All2AllOp = _get_flux_member("All2AllOp")

# inplace cast
inplace_cast_fp32_to_bf16 = _get_flux_member("inplace_cast_fp32_to_bf16")
InplaceCast = _get_flux_member("InplaceCast")

# quantization
Quantization = _get_flux_member("Quantization")

# All2AllTranspose
A2ARingMode = _get_flux_member("A2ARingMode")
AllToAllOption = _get_flux_member("AllToAllOption")

# Gemm + All2AllTranspose
PreAttnAllToAllCommOp = _get_flux_member("PreAttnAllToAllCommOp")
GemmAllToAllTranspose = _get_flux_member("GemmAllToAllTranspose")

# All2AllTranspose + Gemm
AllToAllTransposeGemm = _get_flux_member("AllToAllTransposeGemm")
DisScatterForward = _get_flux_member("DisScatterForward")
DisScatterBackward = _get_flux_member("DisScatterBackward")
All2AllInference = _get_flux_member("All2AllInference")
All2AllSingle = _get_flux_member("All2AllSingle")
AsyncSendRecv = _get_flux_member("AsyncSendRecv")

__all__ = [
    "bsr_reduce",
    "init_flux_shm",
    "create_tensor_list",
    "GroupBarrier",
    "bitwise_check",
    "uniform_initialize",
    "topk_scatter_reduce",
    "load_tuning_record",
    "TuningRecord",
    "ProfilingContext",
    "DistEnvTP",
    "DistEnvTPWithEP",
    "MoeArguments",
    "ReduceScatterOption",
    "RingMode",
    "GemmRS",
    "GemmRSInterNode",
    "AGKernel",
    "AGKernelInterNode",
    "AllToAllTransposeGemm",
    "All2AllOp",
    "GemmOnly",
    "BlockScaleGemm",
    "GemmAllToAllTranspose",
    "GemmGroupedV2",
    "GemmGroupedV3",
    "GemmGroupedV3AGScatter",
    "GemmGroupedV3GatherRS",
    "GemmGroupedV2AGScatterOp",
    "prepare_moe_ag_scatter_args",
    "GemmGroupedV2GatherRSOp",
    "TopkReduceScatterOp",
    "InplaceCast",
    "Quantization",
    "AGRingMode",
    "A2ARingMode",
    "AllGatherOption",
    "AllGatherOp",
    "AllToAllOption",
    "inplace_cast_fp32_to_bf16",
    "get_gemm_rs_threadblock_segments_info",
    "calc_gemm_rs_threadblock_segments_info",
    "calc_scatter_index",
    "calc_moe_triton_blocked_gather_a",
    "DisScatterForward",
    "DisScatterBackward",
    "All2AllInference",
    "PreAttnAllToAllCommOp",
    "All2AllSingle",
    "AsyncSendRecv",
    "_nvshmem_team_split_strided",
    "_nvshmem_team_split_2d",
    "_nvshmem_team_destroy",
    "_nvshmem_team_get_config",
    "_nvshmem_team_translate_pe",
    "_nvshmem_team_npes",
    "_nvshmem_team_my_pe",
]
